#!/usr/bin/python
# -*-coding:utf-8-*-
import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
# plt.switch_backend('agg')

import os
from scipy import stats
import statsmodels.api as sm
import pandas as pd
import numpy as np

# import pickle
# import sys
# import copy
# from time import time
# import datetime
# from datetime import timedelta
# import re

from zg_data_process.zg_data_process import DataProcess

try:
    try:
        from zg_factor_analysis_module.base.FactorAnalysisLibDerivative import FactorAnalysisLibDerivative
        from zg_factor_analysis_module.base.FactorAnalysisLibDerivative import FactorAnalysisLib
        from zg_factor_analysis_module.base.dir_info import factor_lib_dir
    except:
        from .base.FactorAnalysisLibDerivative import FactorAnalysisLibDerivative
        from .base.FactorAnalysisLibDerivative import FactorAnalysisLib
        from .base.dir_info import factor_lib_dir

except:
    from base.FactorAnalysisLib_derivative import FactorAnalysisLib_derivative
    from base.FactorAnalysisLib_derivative import FactorAnalysisLib
    from base.dir_info import factor_lib_dir, basic_data_dir, advance_dir


class FactorAnalysis(object):
    def __init__(self, factor_lib_dir=factor_lib_dir, db='validation'):
        # self.result_save_dir = result_save_dir
        # 初始化数据读取和处理的API
        self.data_reader = FactorAnalysisLib(factor_lib_dir=factor_lib_dir, db=db)

        self.data_process = DataProcess()

    ################################################################################
    # auxillary functions
    # TODO - 添加分组标签
    def _add_group_lebel(self, df, factor_name, method='cap', quantile=5, neutral_quantile=5):
        df = df.reset_index('date').copy()

        threshold = quantile * 20

        if df.shape[0] < threshold:
            return pd.DataFrame()

        if method == 'cap':
            grouping = pd.qcut(df['scale_total_market_size'].rank(method='first'), neutral_quantile, labels=False)
            grouped = df[factor_name].rank(method='first').groupby(grouping).apply(pd.qcut,
                                                                                   quantile,
                                                                                   labels=False) + 1
        elif method == 'sw1':
            grouping = df['sw1_name']
            grouped = df[factor_name].rank(method='first').groupby(grouping).apply(pd.qcut,
                                                                                   quantile,
                                                                                   labels=False) + 1
        elif method == 'ci1':
            grouping = df['ci1_name']
            grouped = df[factor_name].rank(method='first').groupby(grouping).apply(pd.qcut,
                                                                                   quantile,
                                                                                   labels=False) + 1
        else:
            grouped = pd.qcut(df[factor_name].rank(method='first'), quantile, labels=False) + 1

        return grouped

    # TODO - 计算股票权重
    def _cal_stock_weight(self, df, method='ew'):
        if method == 'cir_cap':
            weight = df.groupby('group').apply(
                lambda x: x['scale_circulate_market_size'] / x['scale_circulate_market_size'].sum())
        elif method == 'total_cap':
            weight = df.groupby('group').apply(
                lambda x: x['scale_total_market_size'] / x['scale_total_market_size'].sum())
        else:
            df['weight'] = 1
            weight = df.groupby('group').apply(lambda x: x['weight'] / x.shape[0])

        weight.index = weight.index.droplevel('group')
        weight = weight.sort_index()

        return weight

    # TODO - 计算一些收益指标
    def _calc_annual_return(self, daily_return):
        cum_return = (1 + daily_return).cumprod()

        annual_return = cum_return.ix[-1] ** (243 * 1. / len(cum_return)) - 1

        return annual_return

    def _calc_annual_volatility(self, daily_return):
        return daily_return.std() * np.sqrt(243)

    def _calc_sharpe_ratio(self, daily_return):
        annual_return = self._calc_annual_return(daily_return)
        annual_vol = self._calc_annual_volatility(daily_return)
        sharpe = annual_return / annual_vol
        return sharpe

    def _calc_win_rate(self, daily_return):
        win_rate = len(daily_return[daily_return > 0]) / len(daily_return)

        return win_rate

    def _calc_max_draw_down(self, daily_return):
        cum_return = (1 + daily_return).cumprod()
        cum_peak = np.maximum.accumulate(cum_return)
        dd_to_cum_peak = (cum_peak - cum_return) / cum_peak
        max_dd_end = np.argmax(dd_to_cum_peak)
        max_dd = dd_to_cum_peak[max_dd_end]

        return max_dd

    # TODO 计算IC统计量
    def _cal_ic_statis(self, ic):
        ic_stats = {}

        ic_stats['ic_mean'] = ic.dropna().mean()

        ic_stats['ic_std'] = ic.dropna().std()

        ic_stats['ic_ir'] = ic_stats['ic_mean'] / ic_stats['ic_std']

        ic_stats['pos_ic_ratio'] = ic[ic > 0].shape[0] / ic.shape[0]

        return ic_stats

    ################################################################################

    # TODO - 回归因子收益率计算(日）
    def cal_reg_factor_return(self,
                              label_factor_data,
                              factor_name,
                              ind='sw1',
                              method='OLS'):
        '''
        :param label_factor_data: 某一天的横截面数据，包括：因子值、总市值、流通市值、行业标签、个股收益率等数据项
        :param ind: 选择的行业分类标签 , 可选 sw1(申万一级行业分类)、ci1(中信一级行业分类)
        :param method: 计算因子收益率的方式，默认为OLS，可选WLS
        :return: WLS回归的因子收益率、t统计量 与p值
        '''
        reg_result = {}
        reg_result['factor_return'] = np.nan
        reg_result['tvalues'] = np.nan
        reg_result['pvalues'] = np.nan

        try:
            data = label_factor_data.copy()

            # TODO 获取行业标签
            if ind == 'sw1':
                ind_data = pd.get_dummies(data['sw1_name'])
            else:
                ind_data = pd.get_dummies(data['ci1_name'])

            data[ind_data.columns] = ind_data

            reg_data = data[ind_data.columns.tolist() +
                            [factor_name, 'scale_circulate_market_size', 'scale_total_market_size', 'quote_rate']]

            # TODO 加权最小二乘回归法计算因子收益率，权重为流通市值的平方根
            if method == 'WLS':
                reg = sm.WLS(reg_data['quote_rate'],
                             sm.add_constant(reg_data[['scale_total_market_size', factor_name] + ind_data.columns.tolist()[1:]]),
                             weights=np.sqrt(reg_data['scale_circulate_market_size'])).fit()
            else:
                reg = sm.OLS(reg_data['quote_rate'],
                             sm.add_constant(reg_data[['scale_total_market_size', factor_name] + ind_data.columns.tolist()[1:]])
                             ).fit()

            reg_result['factor_return'] = reg.params[factor_name]

            reg_result['tvalues'] = reg.tvalues[factor_name]

            reg_result['pvalues'] = reg.pvalues[factor_name]

        except Exception as e:
            print(e)
            pass

        return reg_result

    # TODO - 区间回归统计量计算
    def cal_reg_stats(self, reg_result):
        reg_stats = {}

        reg_stats['factor_return_mean'] = reg_result['factor_return'].mean()

        reg_stats['factor_return_std'] = reg_result['factor_return'].std()

        reg_stats['tvalues_mean'] = reg_result['tvalues'].mean()

        reg_stats['pvalues_mean'] = reg_result['pvalues'].mean()

        reg_stats['abs_tvalues_mean'] = reg_result['tvalues'].abs().mean()

        reg_stats['abs_tvalues>2_ratio'] = \
            reg_result[reg_result['tvalues'].abs() >= 2].shape[0] / reg_result.shape[0]

        reg_stats['Fama_MacBeth_tvalues'] = \
            reg_result['factor_return'].mean() / reg_result['factor_return'].std() * np.sqrt(reg_result.shape[0] - 1)

        return reg_stats

    # TODO - 计算因子回归统计分析
    def factor_reg_return(self,
                          label_factor_data,
                          factor_name,
                          industry_type='sw1',
                          method='OLS'):
        """
        :param label_factor_data:  concat_data 为经过处理后的合并数据，包括因子（多个）、行业、总市值、流通市值、个股收益率等数据

        :return:
        每天的因子收益率: dataframe, with date index
        统计量: dict
        """

        cal_data = label_factor_data.set_index(['date', 'stock_code']).copy()

        # reg_items = [factor_name, 'scale_circulate_market_size', 'scale_total_market_size', 'sw1_name']
        reg_items = [factor_name, 'scale_circulate_market_size', 'scale_total_market_size', industry_type+'_name']

        for item in reg_items:
            # print(item)
            cal_data[item] = cal_data.groupby('stock_code')[item].shift(1)

        cal_data = cal_data.dropna()
        total_result = {}

        for index, group in cal_data.groupby('date'):
            # print(index)
            total_result[index] = self.cal_reg_factor_return(group, factor_name, industry_type, method)

        reg_result = pd.DataFrame(total_result).T

        reg_stats = self.cal_reg_stats(reg_result)

        reg_result.index.name = 'date'

        return reg_result, reg_stats

    # TODO - 计算因子IC
    def cal_factor_ic(self, label_factor_data_with_return, factor_name, periods=21, method='rank'):
        '''

        :param label_factor_data_with_return: 因子值和未来收益率的DataFrame
        :param factor_name: 因子名称
        :param periods: 未来预测的时间窗口
        :param method: 方法,: rank, weight

        :return:
        IC: DataFrame
        IC统计量: dict
        '''
        cal_data = label_factor_data_with_return.set_index(['date', 'stock_code']).copy()

        mr = cal_data.groupby('code')['quote_rate'].rolling(periods).sum()
        mr.index = mr.index.droplevel('code')

        cal_data['monthly_return'] = mr
        cal_data[factor_name] = cal_data.groupby('stock_code')[factor_name].shift(periods+1)
        cal_data = cal_data.dropna()

        # TODO 计算加权IC
        def weight_ic(data):
            temp = data.copy()

            temp['weight'] = temp[factor_name].rank().apply(lambda x: 2 ** (2 * x / temp.shape[0] - 2))

            temp['weight_x'] = temp[factor_name] * temp['weight']

            temp['weight_r'] = temp['monthly_return'] * temp['weight']

            return stats.spearmanr(temp['weight_x'], temp['weight_r'])[0]

        if method == 'rank':
            ic = cal_data.groupby('date').apply(lambda x: stats.spearmanr(x[factor_name], x['monthly_return'])[0])
        elif method == 'weight':
            ic = cal_data.groupby('date').apply(lambda x: weight_ic(x))
        else:
            ic = cal_data.groupby('date').apply(lambda x: stats.pearsonr(x[factor_name], x['monthly_return'])[0])

        ic.name = factor_name

        ic_stats = self._cal_ic_statis(ic)

        ic = pd.DataFrame(ic)
        ic.columns = ['ic']

        return ic, ic_stats


    # TODO 计算因子分组统计数据
    def cal_quantile_return(self,
                            label_factor_data,
                            benchmark_index_ret,
                            factor_name,
                            rebalance_date,
                            group_quantile=5,
                            neutral_group_quantile=5,
                            group_method='cap',
                            weight_method='ew'):
        '''
        :param label_factor_data: 带标签的因子数
        :param benchmark_index_ret: 基准指数的收益率
        :param factor_name: 因子名称
        :param rebalance_date: 调仓日数据, list
        :param group_quantile: 分组数
        :param group_method: 分组方法, sw1, ci1, cap, other(不中性化)
        :param weight_method: 配权方法, ew, cir_cap, total_cap

        :return:
        abs_ret: DataFrame, 绝对收益
        excess_ret: DataFrame, 超额收益
        cum_abs_ret: DataFrame, 累计绝对收益
        cum_excess_ret: DataFrame, 累计收益
        position_data: DataFrame, 持仓数据
        quantile_result_stats: DataFrame, 分组统计量
        '''

        # # TODO 数据过滤,若caocat_data已经过滤则不需要这一步
        label_factor_data = label_factor_data.set_index(['date', 'stock_code']).copy()
        cal_data = label_factor_data.copy()
        trade_date = sorted(cal_data.index.get_level_values('date').unique().tolist())

        # TODO - 分组
        cal_data = cal_data.reset_index()
        cal_data = cal_data[cal_data['date'].isin(rebalance_date)]
        cal_data = cal_data.set_index(['date', 'stock_code'])

        grouped_data = cal_data.groupby('date').apply(
            lambda x: self._add_group_lebel(x,
                                            factor_name,
                                            method=group_method,
                                            quantile=group_quantile,
                                            neutral_quantile=neutral_group_quantile)
        )
        cal_data['group'] = grouped_data.loc[cal_data.index]

        cal_data['dt'] = cal_data.index.get_level_values('date')
        w_data = cal_data.dropna()
        weights = w_data.groupby('dt').apply(lambda x: self._cal_stock_weight(x, method=weight_method))
        weights.index = weights.index.droplevel('dt')
        cal_data['weight'] = weights.loc[cal_data.index]

        # TODO - 获取调仓记录数据
        position_data = cal_data.reset_index('stock_code')
        position_data = position_data.loc[rebalance_date]
        position_data = position_data.reset_index()
        position_data = position_data[['date', 'stock_code', 'group', 'weight']]
        position_data = position_data.set_index(['date', 'stock_code'])
        # position_data = position_data.sort_values(['date', 'group'])

        # TODO - 填充每日股票池
        daily_cal_data = cal_data.reset_index().copy()
        # daily_cal_data['date'].unique()

        daily_cal_data = daily_cal_data[daily_cal_data['paused'] == 0.0]
        daily_cal_data = daily_cal_data[daily_cal_data['up_one_line'] == 0.0]
        daily_cal_data = daily_cal_data[daily_cal_data['down_one_line'] == 0.0]
        daily_cal_data = daily_cal_data[['date', 'stock_code', 'group', 'weight']]

        cal_data_list = []
        current_data = pd.DataFrame()
        for td in trade_date:
            # TODO - 选股
            temp = daily_cal_data[daily_cal_data['date'] == td].copy()

            # TODO - 第二天开始计算收益率
            if current_data.shape[0] > 0:
                current_data['date'] = td
                current_data = current_data.set_index(['date', 'stock_code'])

                # TODO - 计算每天收益率
                current_data['quote_rate'] = self.ret_data.loc[current_data.index].copy()
                current_data['quote_rate'] = current_data['quote_rate'].fillna(0.0)

                current_data['weight_return'] = current_data['weight'] * current_data['quote_rate']
                cal_data_list.append(current_data.copy())

                # TODO - 权重变动
                current_data['weight'] = current_data['weight'] * (1+current_data['quote_rate'])

                group_sum = current_data.groupby('group')['weight'].sum()

                current_data['group_sum'] = group_sum.loc[current_data['group']].values
                current_data['weight'] = current_data['weight'] / current_data['group_sum']  # 归一化

                current_data = current_data.reset_index()

            # TODO - 选股当天不计算收益率
            if temp.shape[0] > 0:
                current_data = temp.copy()
                # print(td, 'no return')

        daily_cal_data = pd.concat(cal_data_list, axis=0)
        daily_cal_data = daily_cal_data.drop('group_sum', axis=1)

        cal_data = daily_cal_data.copy()
        del weights, grouped_data, cal_data_list, daily_cal_data

        # TODO 计算加权收益率
        # cal_data['weight_return'] = cal_data['weight'] * cal_data['quote_rate']
        # cal_data['weight_excess_return'] = cal_data['weight'] * cal_data['excess_return']

        # TODO 计算分组收益率
        abs_ret = cal_data.groupby(['date', 'group'])['weight_return'].sum()
        abs_ret = abs_ret.reset_index().pivot(index='date', columns='group', values='weight_return')
        abs_ret['top-bot'] = abs_ret[group_quantile] - abs_ret[1]
        abs_ret['bot-top'] = abs_ret[1] - abs_ret[group_quantile]

        # index_ret = self.get_daily_index_trade_data(bench)

        abs_ret['bench'] = benchmark_index_ret.loc[abs_ret.index, 'quote_rate']
        cum_abs_ret = (1 + abs_ret).cumprod()

        excess_ret = abs_ret.iloc[:, :group_quantile].apply(lambda x: x - abs_ret['bench'])
        cum_excess_ret = (1 + excess_ret).cumprod()

        ####
        quantile_result_stats = {}
        quantile_result_stats['annual_return'] = abs_ret.apply(lambda x: self._calc_annual_return(x))
        quantile_result_stats['annual_vol'] = abs_ret.apply(lambda x: self._calc_annual_volatility(x))
        quantile_result_stats['SR'] = abs_ret.apply(lambda x: self._calc_sharpe_ratio(x))
        quantile_result_stats['winrate'] = abs_ret.apply(lambda x: self._calc_win_rate(x))
        quantile_result_stats['max_draw'] = abs_ret.apply(lambda x: self._calc_max_draw_down(x))

        quantile_result_stats['excess_annual_return'] = excess_ret.apply(lambda x: self._calc_annual_return(x))
        quantile_result_stats['TE'] = excess_ret.apply(lambda x: self._calc_annual_volatility(x))
        quantile_result_stats['IR'] = excess_ret.apply(lambda x: self._calc_sharpe_ratio(x))
        quantile_result_stats['excess_winrate'] = excess_ret.apply(lambda x: self._calc_win_rate(x))
        quantile_result_stats['excess_max_draw'] = excess_ret.apply(lambda x: self._calc_max_draw_down(x))

        # TODO - 输出格式化
        abs_ret.columns = ['g%d' % (g+1) for g in range(group_quantile)] + abs_ret.columns[group_quantile:].tolist()
        excess_ret.columns = ['g%d' % (g+1) for g in range(group_quantile)] + excess_ret.columns[group_quantile:].tolist()
        cum_abs_ret.columns = ['g%d' % (g+1) for g in range(group_quantile)] + cum_abs_ret.columns[group_quantile:].tolist()
        cum_excess_ret.columns = ['g%d' % (g+1) for g in range(group_quantile)] + cum_excess_ret.columns[group_quantile:].tolist()

        quantile_result_stats = pd.DataFrame(quantile_result_stats)
        quantile_result_stats.index.name = 'statstics'

        return abs_ret, excess_ret, cum_abs_ret, cum_excess_ret, position_data, quantile_result_stats

    #TODO 画分组绝对收益曲线图
    def plot_periods_abs_ret(self, ret, factor_name, plot_save_dir, benchmark=None):
        cum_ret=(ret+1).cumprod()

        # plt.figure(figsize=(40, 20))
        # cum_ret.plot(legend=True,fontsize=8, lw=1, alpha=0.6)
        cum_ret.plot(legend=True, fontsize=12, lw=1.5, alpha=0.8, figsize=(20, 10))
        plt.title(factor_name +' quantile return')
        plt.legend(loc=2, bbox_to_anchor=(1, 1), borderaxespad=0, fontsize=12)

        if benchmark is None:
            save_filename = '%s_quantile_abs_return.png' % (factor_name)
        else:
            save_filename = '%s_%s_pool_quantile_abs_return.png' % (factor_name, benchmark)

        plt.savefig(os.path.join(plot_save_dir, save_filename), dpi=200)
        plt.close()

    # TODO 画分组超额收益曲线图
    def plot_periods_excess_ret(self, ret, factor_name, plot_save_dir, benchmark=None):
        cum_ret = (ret+1).cumprod()

        # plt.figure(figsize=(40, 20))
        # cum_ret.plot(legend=True,fontsize=8, lw=1, alpha=0.6)
        cum_ret.plot(legend=True,fontsize=12, lw=1.5, alpha=0.8, figsize=(20, 10))
        plt.title(factor_name +' quantile excess return')
        plt.legend(loc=2, bbox_to_anchor=(1, 1), borderaxespad=0, fontsize=12)

        if benchmark is None:
            save_filename = '%s_quantile_excess_return.png' % (factor_name)
        else:
            save_filename = '%s_%s_pool_quantile_excess_return.png' % (factor_name, benchmark)

        plt.savefig(os.path.join(plot_save_dir, save_filename), dpi=200)
        plt.close()

    #TODO 计算因子回归收益率的相关性——在计算出回归因子收益率的基础上
    def cal_reg_factor_return_corr(self, factor1_ret, factor2_ret):
        # factor1_ret = self.get_factor_reg_return(factor1,pool)
        # factor2_ret = self.get_factor_reg_return(factor2,pool)
        df = pd.concat([factor1_ret,
                        factor2_ret],
                       axis=1)

        corr = df.corr().iloc[0,1]

        return corr

    #TODO 计算因子IC的相关性——在计算出IC的基础上
    def cal_factor_ic_corr(self, factor1_IC, factor2_IC):
        # factor1_IC = self.get_factor_ic_data(factor1,pool)
        # factor2_IC = self.get_factor_ic_data(factor2,pool)
        df = pd.concat([factor1_IC, factor2_IC],axis=1)

        corr=df.corr().iloc[0,1]

        return corr

    # TODO - 计算因子的选股交叉度——在分层回测的持仓数据基础之上
    def cal_factors_stock_overlap(self, factor1_position, factor2_position):
        try:
            # factor1_position = self.get_factor_position_data(factor1, pool=pool, rebalance=rebalance, group=group, weight=weight)
            # factor2_position = self.get_factor_position_data(factor2, pool=pool, rebalance=rebalance, group=group, weight=weight)
            com_date_list = sorted(
                list(factor1_position.index.get_level_values('date').unique().intersection(factor2_position.index.get_level_values('date').unique()))
            )

            temp = {}

            for date in com_date_list:
                df1 = factor1_position.xs(date)
                df2 = factor2_position.xs(date)
                overlap_ratio = {}
                for q in df1.index.unique():
                    stocks1 = df1.loc[q]['stock_code'].tolist()
                    stocks2 = df2.loc[q]['stock_code'].tolist()
                    overlap_ratio[q] = (len(set(stocks1).intersection(set(stocks2))) / max(len(stocks1), len(stocks2)))
                overlap_ratio['mean'] = pd.Series(overlap_ratio).mean()
                temp[date] = overlap_ratio

            return pd.DataFrame(temp).T
        except Exception as e:
            print(e)
            return np.nan

if __name__ == '__main__':
    pass



